import sqlite3 
from structs import *
from typing import List
class Recorder:
    def __init__(
            self,
            db_file:str,
            table_names:str|List[str]
    ):
        self.conn=sqlite3.connect(db_file)
        self.conn.row_factory = sqlite3.Row
        self.cur = self.conn.cursor()
        if isinstance(table_names,List):
            for table_name in table_names:
                self._create_table(table_name)
        else:
            self._create_table(table_names)

    def _create_table(self,table_name:str):
        print(f"table name is {table_name}")
        self.cur.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                tp_world_size INTEGER,
                batch_size INTEGER,
                input_len INTEGER,
                output_len INTEGER,
                avg_prefill_time_usage REAL,
                avg_decoding_time_usage REAL,
                prefill_time_stddev REAL,
                decoding_time_stddev REAL,
                generation_throughput REAL
            )
            """
        )
        self.conn.commit()


    def insert_record(
            self,
            worker_param:WorkerParam,
            input_param:InputParam,
             avg_prefill_time_usage:float,
                avg_decoding_time_usage:float,
                prefill_time_stddev:float,
                decoding_time_stddev:float,
            generation_throughput:float
    ):
        table_name=worker_param.model_name
        self.cur.execute(
                f"""
                INSERT INTO {table_name} (tp_world_size, batch_size, input_len, output_len, avg_prefill_time_usage, avg_decoding_time_usage, prefill_time_stddev, decoding_time_stddev,generation_throughput) VALUES (?, ?, ?, ?, ?, ?, ?, ?,?)
                """,
                (worker_param.tp_world_size,
                 input_param.batch_size, input_param.input_len, input_param.output_len,
                 avg_prefill_time_usage, avg_decoding_time_usage, prefill_time_stddev, decoding_time_stddev,
                 generation_throughput)
            )
        self.conn.commit()
    
    def query_record(
        self,
        worker_param: WorkerParam,
        input_param: InputParam
    ) -> sqlite3.Row:
        """
        Query the record from the database. If the record does not exist, return None
        """
        model_name=worker_param.model_name
        self.cur.execute(
            f"""
            SELECT * FROM {model_name}
            WHERE tp_world_size = ? AND
                batch_size = ? AND input_len = ? AND output_len = ?
            """,
            (worker_param.tp_world_size,
             input_param.batch_size, input_param.input_len, input_param.output_len)
        )
        return self.cur.fetchone()